#!/usr/bin/python

# TODO:
# -- add rel-based selection for training?

import os,time
import numpy
from collections import Counter
from pylab import *
import tables,scipy
from scipy import fftpack
from scipy.cluster import vq
from scipy.ndimage import morphology,filters,interpolation
from ocrolib import docproc,improc
import pyflann
import random
import tables
import multiprocessing
import fcntl
from pdb import pm
from collections import Counter
from llpy import mlp
import cPickle
from llpy import variants
from llpy.sutils import *

nnet = mlp.CAutoMLP()
mlp.maxthreads.value = multiprocessing.cpu_count()
mlp.maxthreads_train.value = multiprocessing.cpu_count()

import argparse
parser = argparse.ArgumentParser(description = "Train neural network character models.")
parser.add_argument('data',help="data file")
parser.add_argument('-o','--output',default="output.mlp",help="output file")
parser.add_argument('-N','--nsamples',default=8000000,type=int,help="number of training samples")
parser.add_argument('-r','--rounds',default=48,type=int,help="number of rounds")
parser.add_argument("-c",'--noreject',action="store_true",help="do not use reject samples for training")
parser.add_argument("-g","--nogeo",action="store_true",help="turn off geometry")
#cmdline = "-N 99999999 -r 10 --reject DATA/all_training.h5".split()
#args = parser.parse_args(cmdline)
args = parser.parse_args()
clist = None

with tables.openFile(args.data) as db:
    nsamples = len(db.root.classes)
    ndims = db.root.patches[0].size
    classes = array(db.root.classes[:nsamples])
    ignore = sum(classes==ord("_"))
    if ignore>0:
        print "warning:",ignore,"ignorable samples"

    if args.noreject:
        print "NOT training rejects classes"
        indexes = [i for i,c in enumerate(classes) if c not in [ord("_"),ord("~")]]
    else:
        print "training rejects classes"
        indexes = [i for i,c in enumerate(classes) if c!=ord("_")]

    if args.nsamples<len(indexes):
        indexes = sorted(random.sample(indexes,args.nsamples))

    print "quick check whether samples are in range"
    for i in range(10):
        lo = amin(db.root.patches[i])
        hi = amax(db.root.patches[i])
        assert lo>-0.1
        assert hi>0.9
        assert hi<1.1
    print "seems OK"

    print db
    images = (db.root.patches[i] for i in indexes)
    chars = [udecode(classes[i]) for i in indexes]

    print "#samples",len(indexes)
    print "#classes",len(set(chars))

    if "rel" in db.root and not args.nogeo:
        print "using geometry information"
        rels = (db.root.rel[i] for i in indexes)
    else:
        print "no geometry information"
        rels = None

    nnet.max_rounds = args.rounds
    nnet.ctrain(images,chars,rels,verbose=1)
    print nnet.classlabels

with open(args.output,"w") as stream:
    cPickle.dump(nnet,stream,2)
